#!/usr/bin/env python
from numpy import *
import sys

class Data:
  def __init__(self):
    self._var2name = dict();
    self._name2var = dict();
    
    self._sample2name = dict();
    self._name2sample = dict();
  
  def setVariableName(self,var_id,var_name):
    var_id = str(var_id);
    var_name = str(var_name);
    self._var2name[var_id] = var_name;
    self._name2var[var_name] = var_id;
  
  def getVariableName(self,var_id):
    var_id = str(var_id);
    if var_id in self._var2name:
      return self._var2name[var_id];
    else:
      return str(var_id);
  
  def setSampleName(self,sample_id,sample_name):
    sample_id = str(sample_id);
    sample_name = str(sample_name);
    self._sample2name[sample_id] = sample_name;
    self._name2sample[sample_name] = sample_id;
  
  def getSampleName(self,sample_id):
    sample_id = str(sample_id);
    if sample_id in self._sample2name:
      return self._sample2name[sample_id];
    else:
      return str(sample_id);
  
  def setValue(self,row,col,value):
    raise NotImplementedError;
  
  def getValue(self,row,col):
    raise NotImplementedError;

class DiscreteData(Data):
  def __init__(self,nrows,ncols):
    Data.__init__(self);
    self._data = zeros((nrows,ncols),dtype=int8);
  
  def setValue(self,row,col,value):
    self._data[row][col] = value;
  
  def getValue(self,row,col):
    return self._data[row][col];
  
  def numRows(self):
    return shape(self._data)[0];
  
  def numCols(self):
    return shape(self._data)[1];

# Wrapper class (well not exactly a wrapper, but whatever) for DiscreteData.
# Basically it translates rows and columns into variables and observations. 
class BayesData(DiscreteData):
  def __init__(self,numVars,numObs):
    DiscreteData.__init__(self,numObs,numVars);
    self._sm = StateManager(numVars);
    
  def setValue(self,variable,observation,value):
    intValue = self._sm.getInt(variable,value);
    DiscreteData.setValue(self,observation,variable,intValue);
    
  def getValue(self,variable,observation):
    return DiscreteData.getValue(self,observation,variable);
  
  def setClassVar(self,classvar):
    self._class = int(classvar);
    
  def getClassVar(self,classvar):
    return self._class;
  
  def numVars(self):
    return self.numCols();
  
  def numObs(self):
    return self.numRows();
  
  def write(self,out=sys.stdout):
    # Print header.
    out.write(self.getVariableName(0));
    for i in xrange(1,self.numCols()):
      out.write("\t" + self.getVariableName(i));
      
    out.write("\n");
    
    # Print data.
    for i in xrange(self.numRows()):
      out.write('\t'.join([str(e) for e in self._data[i]]));
      out.write("\n");

class StateManager:
  def __init__(self,numVars):
    self._int2str = [dict() for i in xrange(numVars)];
    self._str2int = [dict() for i in xrange(numVars)];
    self._nextAvailableState = [0]*numVars;
    
  def getInt(self,var,state):
    state = str(state);
    var = int(var);
    
    value = self._str2int[var].get(state,None);
    if value == None:
      # Never seen this state before, map it.
      value = self._nextAvailableState[var];
      self._str2int[var][state] = self._nextAvailableState[var];
      self._int2str[var][self._nextAvailableState[var]] = state;
      self._nextAvailableState[var] += 1;
      return value;
    else:
      # We've seen it before, just return it. 
      return value;
  
  def getString(self,var,state):
    state = int(state);
    var = int(var);
    
    value = self._int2str[var].get(state,None);
    if value == None:
      raise ValueError, "Conversion from state to string does not exist."
    else:
      return value;
  